# Copyright (c) Facebook, Inc. and its affiliates.
# Copied from: https://github.com/facebookresearch/detectron2/blob/master/demo/predictor.py
import atexit
import bisect
import multiprocessing as mp
from collections import deque

import cv2
import numpy as np
import torch
from detectron2.data import MetadataCatalog
from detectron2.engine.defaults import DefaultPredictor
from detectron2.utils.video_visualizer import VideoVisualizer
from detectron2.utils.visualizer import ColorMode, Visualizer


class VisualizationDemo(object):
    def __init__(self, cfg, instance_mode=ColorMode.IMAGE, parallel=False):
        """
        Args:
            cfg (CfgNode):
            instance_mode (ColorMode):
            parallel (bool): whether to run the model in different processes from visualization.
                Useful since the visualization logic can be slow.
        """
        self.metadata = MetadataCatalog.get(
            cfg.DATASETS.TEST[0] if len(cfg.DATASETS.TEST) else "__unused"
        )
        self.cpu_device = torch.device("cpu")
        self.instance_mode = instance_mode

        self.parallel = parallel
        if parallel:
            num_gpu = torch.cuda.device_count()
            self.predictor = AsyncPredictor(cfg, num_gpus=num_gpu)
        else:
            self.predictor = DefaultPredictor(cfg)

    def run_on_image(self, image):
        """
        Args:
            image (np.ndarray): an image of shape (H, W, C) (in BGR order).
                This is the format used by OpenCV.
        Returns:
            predictions (dict): the output of the model.
            vis_output (VisImage): the visualized image output.
        """
        vis_output = None
        predictions = self.predictor(image)
        # Convert image from OpenCV BGR format to Matplotlib RGB format.
        image = image[:, :, ::-1]
        visualizer = Visualizer(image, self.metadata, instance_mode=self.instance_mode)
        if "panoptic_seg" in predictions:
            panoptic_seg, segments_info = predictions["panoptic_seg"]
            vis_output = visualizer.draw_panoptic_seg_predictions(
                panoptic_seg.to(self.cpu_device), segments_info
            )
        else:
            if "sem_seg" in predictions:
                vis_output = visualizer.draw_sem_seg(
                    predictions["sem_seg"].argmax(dim=0).to(self.cpu_device)
                )
            if "instances" in predictions:
                instances = predictions["instances"].to(self.cpu_device)
                vis_output = visualizer.draw_instance_predictions(predictions=instances)

        return predictions, vis_output
    
    def run_on_dianjiao(self, image, threshold=0.5):
        """
        Args:
            image (np.ndarray): an image of shape (H, W, C) (in BGR order).
                This is the format used by OpenCV.
        Returns:
            predictions (dict): the output of the model.
            vis_output (VisImage): the visualized image output.
        """
        vis_output = None
        predictions = self.predictor(image)
        # Convert image from OpenCV BGR format to Matplotlib RGB format.
        image = image[:, :, ::-1]
        visualizer = Visualizer(image, self.metadata, instance_mode=self.instance_mode)
        assert len(predictions["sem_seg"]) == 1, "dianjiao only support binary classification."
        mask = predictions["sem_seg"][0]
        mask[mask > threshold] = 255
        mask[mask < threshold] = 0
        # 此时的mask已经是分割预测结果
        # 将像素值为0的像素改为255，将像素值为255的像素改为0
        draw_mask = 255 - mask.type(torch.int64)
        vis_output = visualizer.draw_dianjiao_seg(draw_mask.to(self.cpu_device), alpha=0.4)

        return predictions, vis_output
    
    def run_model(self, image, threshold=0.5, use_sigmoid=False):
        predictions = self.predictor(image)
        # Convert image from OpenCV BGR format to Matplotlib RGB format.
        assert len(predictions["sem_seg"]) == 1, "dianjiao only support binary classification."
        mask = predictions["sem_seg"][0]
        if use_sigmoid:
            mask = torch.sigmoid(mask)
        mask[mask > threshold] = 255
        mask[mask < threshold] = 0
        # 此时的mask已经是分割预测结果, 0为背景，255为前景

        return mask.cpu().numpy()
    
    def run_model_parallel(self, img_list, threshold=0.5, use_sigmoid=False):
        with torch.no_grad():
            batch_inputs = []
            for original_image in img_list:
                if self.predictor.input_format == "RGB":
                    # whether the model expects BGR inputs or RGB
                    original_image = original_image[:, :, ::-1]
                height, width = original_image.shape[:2]
                image = self.predictor.aug.get_transform(original_image).apply_image(original_image)
                image = torch.as_tensor(image.astype("float32").transpose(2, 0, 1))
                batch_inputs.append({"image": image, "height": height, "width": width})
            
            predictions = self.predictor.model(batch_inputs)

        # Convert image from OpenCV BGR format to Matplotlib RGB format.
        assert len(predictions[0]["sem_seg"]) == 1, "dianjiao only support binary classification."

        mask_list = []
        for pred in predictions:
            mask = pred["sem_seg"][0]
            if use_sigmoid:
                mask = torch.sigmoid(mask)
            mask[mask > threshold] = 255
            mask[mask < threshold] = 0
            # 此时的mask已经是分割预测结果, 0为背景，255为前景
            mask_list.append(mask.cpu().numpy())

        return mask_list



class AsyncPredictor:
    """
    A predictor that runs the model asynchronously, possibly on >1 GPUs.
    Because rendering the visualization takes considerably amount of time,
    this helps improve throughput a little bit when rendering videos.
    """

    class _StopToken:
        pass

    class _PredictWorker(mp.Process):
        def __init__(self, cfg, task_queue, result_queue):
            self.cfg = cfg
            self.task_queue = task_queue
            self.result_queue = result_queue
            super().__init__()

        def run(self):
            predictor = DefaultPredictor(self.cfg)

            while True:
                task = self.task_queue.get()
                if isinstance(task, AsyncPredictor._StopToken):
                    break
                idx, data = task
                result = predictor(data)
                self.result_queue.put((idx, result))

    def __init__(self, cfg, num_gpus: int = 1):
        """
        Args:
            cfg (CfgNode):
            num_gpus (int): if 0, will run on CPU
        """
        num_workers = max(num_gpus, 1)
        self.task_queue = mp.Queue(maxsize=num_workers * 3)
        self.result_queue = mp.Queue(maxsize=num_workers * 3)
        self.procs = []
        for gpuid in range(max(num_gpus, 1)):
            cfg = cfg.clone()
            cfg.defrost()
            cfg.MODEL.DEVICE = "cuda:{}".format(gpuid) if num_gpus > 0 else "cpu"
            self.procs.append(
                AsyncPredictor._PredictWorker(cfg, self.task_queue, self.result_queue)
            )

        self.put_idx = 0
        self.get_idx = 0
        self.result_rank = []
        self.result_data = []

        for p in self.procs:
            p.start()
        atexit.register(self.shutdown)

    def put(self, image):
        self.put_idx += 1
        self.task_queue.put((self.put_idx, image))

    def get(self):
        self.get_idx += 1  # the index needed for this request
        if len(self.result_rank) and self.result_rank[0] == self.get_idx:
            res = self.result_data[0]
            del self.result_data[0], self.result_rank[0]
            return res

        while True:
            # make sure the results are returned in the correct order
            idx, res = self.result_queue.get()
            if idx == self.get_idx:
                return res
            insert = bisect.bisect(self.result_rank, idx)
            self.result_rank.insert(insert, idx)
            self.result_data.insert(insert, res)

    def __len__(self):
        return self.put_idx - self.get_idx

    def __call__(self, image):
        self.put(image)
        return self.get()

    def shutdown(self):
        for _ in self.procs:
            self.task_queue.put(AsyncPredictor._StopToken())

    @property
    def default_buffer_size(self):
        return len(self.procs) * 5



def ExtractROI(image, image_roi):
    x,y,h,w = image_roi
    img_shape = image.shape
    if len(img_shape) == 3:
        roi = image[y:y+h, x:x+w, :]
    else:
        roi = image[y:y+h, x:x+w]

    return roi


def Split_6(img, params):

    dict_part2img={}
    for part in ['l','r','tl','tr','bl','br']:
        roi_coord=params['roi_{}'.format(part)]
        patch=ExtractROI(img, (roi_coord[0],roi_coord[1], roi_coord[3],roi_coord[2],))
        if part =='l' or part=='r':
            patch = np.rot90(patch).copy()
        dict_part2img[part]=patch

    return dict_part2img

def Merge_6(black_img, dict_part2mask,params):

    final_image = black_img.copy()
    for part in ['l','r','tl','tr','bl','br']:
        roi_coord=params['roi_{}'.format(part)]
        # roi_coord[2] = roi_coord[2]-roi_coord[0]
        # roi_coord[3] = roi_coord[3]-roi_coord[1]
        ReplaceROI(final_image, dict_part2mask[part], (roi_coord[0],roi_coord[1], roi_coord[3],roi_coord[2],))

    return final_image

def ReplaceROI(src_image, patch_image, image_roi):
    x,y,h,w = image_roi
    img_shape = src_image.shape
    if len(img_shape) == 3:
        src_image[y:y+h, x:x+w, :] = patch_image
    else:
        src_image[y:y+h, x:x+w] += patch_image

    src_image[y:y+h, x:x+w] = np.where(src_image[y:y+h, x:x+w] > 0, 255, 0)